import numpy as np
import math
from matplotlib import pyplot as plt
import cv2
class Image(object):
def __init__(self, path):
self.path = path
self.rgb_image = None
self.bgr_image = None
self.gray_image = None
def read_image(self, return_image = False):
self.rgb_image = plt.imread(self.path)
if return_image:
return self.rgb_image
def bgr(self, return_image = False):
self.bgr_image = np.flip(plt.imread(self.path), 2)
if return_image:
return self.bgr_image
def gray(self, return_image = False):
self.gray_image = cv2.cvtColor(plt.imread(self.path), cv2.COLOR_RGB2GRAY)
if return_image:
return self.gray_imagezoom_img
@staticmethod
def show(image, title = 'image'):
if len(image.shape) == 3:
plt.imshow(image)
else:
plt.imshow(image, cmap = 'gray')
plt.title(title)
@staticmethod
def show_all(image_list, title_list):
assert len(image_list) == len(title_list), "Incompatible lengths of lists!"
N = len(image_list)
cols = 5
rows = N // cols
plt.figure(figsize=[50, 50])
for i in range(rows):
for j in range(cols):
plt.subplot(rows, cols, i * cols + j + 1)
Image.show(image_list[i * cols + j], title_list[i * cols + j])
plt.show()
class Augmentation:
@staticmethod
def flip_img(img, axis):
return np.flip(img, axis)
@staticmethod
def rotate_img(img, theta):
center = tuple(np.array(img.shape[1::-1]) / 2)
rotated_matrix = cv2.getRotationMatrix2D(center, theta, 1.0)
rotated_img = cv2.warpAffine(img, rotated_matrix, img.shape[1::-1], flags=cv2.INTER_LINEAR)
return rotated_img
@staticmethod
def crop_img(img, y, x, index_coordinate):
if index_coordinate[0] + y >= img.shape[0]:
y = img.shape[0] - index_coordinate[0] - 1
if index_coordinate[1] + x >= img.shape[1]:
y = img.shape[1] - index_coordinate[1] - 1
return img[index_coordinate[0]:index_coordinate[0] + y, index_coordinate[1]:index_coordinate[1] + x]
@staticmethod
def gaussian_blur(img, kernel_size, sigma):
rand_size = np.random.choice([3, 5, 7, 9 ,11], 1, [1/5] * 5)
blurred_img = cv2.GaussianBlur(img, kernel_size, sigma, cv2.BORDER_DEFAULT)
return blurred_img
@staticmethod
def median_blur(img, kernel_size):
rand_size = np.random.choice([3, 5, 7, 9 ,11], 1, [1/5] * 5)
blurred_img = cv2.medianBlur(img, kernel_size)
return blurred_img
@staticmethod
def zoom_img(img, zoom_level, pos):
x = int(img.shape[1] * 1/zoom_level)
y = int(img.shape[0] * 1/zoom_level)
dimension = (img.shape[1], img.shape[0])
new_img = Augmentation.crop_img(img, x, y, (pos[0], pos[1]))
zoomed_img = cv2.resize(new_img, dimension, interpolation = cv2.INTER_AREA)
return zoomed_img
class Use_Augmentation (Augmentation):
from enum import Enum
class AugmentationEnum(Enum):
FLIP = 1
ROTATE = 2
CROP = 3
GAUSSIAN_BLUR = 4
MEDIAN_BLUR = 5
ZOOM = 6
@staticmethod
def random_augmentations(img, n, m=[1/6, 1/6, 1/6, 1/6, 1/6, 1/6]):
# n = number of augmentations to commit
# m = list of probabilities
rand_size = np.random.choice([3, 5, 7, 9 ,11], 1, [1/5] * 5)
aug_imges = list()
current_img = np.zeros((img.shape))
aug_types = np.random.choice([Use_Augmentation.AugmentationEnum.FLIP,
Use_Augmentation.AugmentationEnum.ROTATE,
Use_Augmentation.AugmentationEnum.CROP,
Use_Augmentation.AugmentationEnum.GAUSSIAN_BLUR,
Use_Augmentation.AugmentationEnum.MEDIAN_BLUR,
Use_Augmentation.AugmentationEnum.ZOOM], n, m)
for i in range(len(aug_types)):
if aug_types[i] == Use_Augmentation.AugmentationEnum.FLIP:
current_img = Augmentation.flip_img(img, np.random.randint(2))
elif aug_types[i] == Use_Augmentation.AugmentationEnum.ROTATE:
current_img = Augmentation.rotate_img(img, np.random.randint(180))
elif aug_types[i] == Use_Augmentation.AugmentationEnum.CROP:
current_img = Augmentation.crop_img(img, np.random.randint(300), np.random.randint(300), (np.random.randint(300), np.random.randint(300)))
elif aug_types[i] == Use_Augmentation.AugmentationEnum.GAUSSIAN_BLUR:
current_img = Augmentation.gaussian_blur(img, (rand_size, rand_size), np.random.randint(50))
rand_size = np.random.choice([3, 5, 7, 9 ,11], 1, [1/5] * 5)
elif aug_types[i] == Use_Augmentation.AugmentationEnum.MEDIAN_BLUR:
current_img = Augmentation.median_blur(img, rand_size)
rand_size = np.random.choice([3, 5, 7, 9 ,11], 1, [1/5] * 5)
elif aug_types[i] == Use_Augmentation.AugmentationEnum.ZOOM:
current_img = Augmentation.zoom_img(img, np.random.randint(1, 4), (np.random.randint(300), np.random.randint(300)))
aug_imges.append(current_img)
Image.show_all(aug_imges, [str(i) for i in range(n)])
stopsign = Image("./stopsign.jpg").read_image(True)
Use_Augmentation.random_augmentations(stopsign, 50)